#!/usr/bin/env python3
"""
批量视频转录+润色系统
- 阶段1: 多进程Whisper Tiny转录
- 阶段2: DeepSeek API并发润色
"""

import json
import asyncio
import time
from pathlib import Path
from multiprocessing import Pool, Manager, cpu_count
from faster_whisper import WhisperModel
from openai import AsyncOpenAI
from tqdm import tqdm
import argparse

# ============= 配置 =============
WHISPER_MODEL = "tiny"
WHISPER_DEVICE = "cpu"
WHISPER_COMPUTE_TYPE = "int8"

# DeepSeek API配置
DEEPSEEK_API_KEY = "your-api-key-here"  # 需要替换为你的API Key
DEEPSEEK_BASE_URL = "https://api.deepseek.com"

OUTPUT_DIR = Path("batch_output")
OUTPUT_DIR.mkdir(exist_ok=True)


# ============= 阶段1: Whisper转录 (多进程) =============

def transcribe_single_video(args):
    """
    转录单个视频（在子进程中执行）

    Args:
        args: (video_path, video_index, total_videos)

    Returns:
        dict: 转录结果
    """
    video_path, video_index, total_videos = args

    try:
        # 在每个进程中加载模型
        model = WhisperModel(
            WHISPER_MODEL,
            device=WHISPER_DEVICE,
            compute_type=WHISPER_COMPUTE_TYPE
        )

        # 转录
        segments, info = model.transcribe(
            video_path,
            language="zh",
            initial_prompt="以下是普通话的句子。"
        )

        # 收集文本
        text_parts = []
        for segment in segments:
            text_parts.append(segment.text.strip())

        full_text = "\n".join(text_parts)

        return {
            "success": True,
            "video_path": video_path,
            "video_name": Path(video_path).name,
            "raw_text": full_text,
            "duration": info.duration,
            "language": info.language
        }

    except Exception as e:
        return {
            "success": False,
            "video_path": video_path,
            "video_name": Path(video_path).name,
            "error": str(e)
        }


def batch_transcribe_multiprocess(video_paths, num_workers=4):
    """
    多进程批量转录

    Args:
        video_paths: 视频路径列表
        num_workers: 进程数（建议2-4）

    Returns:
        list: 转录结果列表
    """
    print(f"\n{'='*80}")
    print(f"阶段1: Whisper转录")
    print(f"{'='*80}")
    print(f"模型: {WHISPER_MODEL}")
    print(f"视频数量: {len(video_paths)}")
    print(f"并发进程: {num_workers}")
    print(f"预估时间: {len(video_paths) / (15 * num_workers):.1f} 分钟")
    print(f"{'='*80}\n")

    # 准备参数
    args_list = [
        (video_path, i+1, len(video_paths))
        for i, video_path in enumerate(video_paths)
    ]

    # 多进程执行
    start_time = time.time()

    with Pool(processes=num_workers) as pool:
        # 使用tqdm显示进度
        results = list(tqdm(
            pool.imap(transcribe_single_video, args_list),
            total=len(video_paths),
            desc="转录进度",
            unit="视频"
        ))

    elapsed = time.time() - start_time

    # 统计
    success_count = sum(1 for r in results if r["success"])
    failed_count = len(results) - success_count

    print(f"\n转录完成！")
    print(f"  成功: {success_count}")
    print(f"  失败: {failed_count}")
    print(f"  耗时: {elapsed/60:.1f} 分钟")
    print(f"  速度: {len(video_paths)/elapsed*60:.1f} 视频/分钟\n")

    return results


# ============= 阶段2: DeepSeek润色 (异步并发) =============

async def polish_with_deepseek(text, client, semaphore):
    """
    使用DeepSeek API润色文本（异步）

    Args:
        text: 原始文本
        client: AsyncOpenAI客户端
        semaphore: 并发控制信号量

    Returns:
        str: 润色后的文本
    """
    async with semaphore:
        try:
            response = await client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {
                        "role": "system",
                        "content": "你是一个专业的文本校对助手。请修正语音识别错误，修正错别字，使文本更加通顺流畅，但保持原意不变。直接输出润色后的文本，不要添加任何解释。"
                    },
                    {
                        "role": "user",
                        "content": f"请润色以下视频转录文本：\n\n{text}"
                    }
                ],
                temperature=0.3
            )

            return response.choices[0].message.content.strip()

        except Exception as e:
            print(f"润色失败: {str(e)}")
            return text  # 失败时返回原文


async def batch_polish_async(transcribe_results, concurrency=20):
    """
    异步批量润色

    Args:
        transcribe_results: 转录结果列表
        concurrency: 并发数

    Returns:
        list: 润色后的结果
    """
    print(f"\n{'='*80}")
    print(f"阶段2: DeepSeek润色")
    print(f"{'='*80}")
    print(f"文本数量: {len(transcribe_results)}")
    print(f"并发请求: {concurrency}")
    print(f"{'='*80}\n")

    # 初始化客户端
    client = AsyncOpenAI(
        api_key=DEEPSEEK_API_KEY,
        base_url=DEEPSEEK_BASE_URL
    )

    # 并发控制
    semaphore = asyncio.Semaphore(concurrency)

    # 只润色成功的转录
    tasks = []
    for result in transcribe_results:
        if result["success"]:
            task = polish_with_deepseek(
                result["raw_text"],
                client,
                semaphore
            )
            tasks.append(task)
        else:
            tasks.append(asyncio.sleep(0))  # 占位

    # 执行
    start_time = time.time()

    polished_texts = []
    for i, task in enumerate(asyncio.as_completed(tasks)):
        polished_text = await task
        polished_texts.append(polished_text)

        # 进度显示
        if (i + 1) % 10 == 0:
            print(f"已润色: {i+1}/{len(tasks)}")

    elapsed = time.time() - start_time

    print(f"\n润色完成！")
    print(f"  耗时: {elapsed:.1f} 秒")
    print(f"  速度: {len(tasks)/elapsed:.1f} 个/秒\n")

    # 合并结果
    final_results = []
    polish_idx = 0
    for result in transcribe_results:
        if result["success"]:
            result["polished_text"] = polished_texts[polish_idx]
            polish_idx += 1
        else:
            result["polished_text"] = None
        final_results.append(result)

    return final_results


# ============= 保存结果 =============

def save_results(results):
    """保存处理结果"""
    print(f"\n{'='*80}")
    print(f"保存结果")
    print(f"{'='*80}\n")

    for result in results:
        if not result["success"]:
            continue

        video_name = Path(result["video_name"]).stem

        # 保存原始转录
        raw_file = OUTPUT_DIR / f"{video_name}_raw.txt"
        raw_file.write_text(result["raw_text"], encoding="utf-8")

        # 保存润色文本
        if result.get("polished_text"):
            polished_file = OUTPUT_DIR / f"{video_name}_polished.txt"
            polished_file.write_text(result["polished_text"], encoding="utf-8")

    # 保存完整JSON
    json_file = OUTPUT_DIR / "batch_results.json"
    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)

    print(f"结果已保存到: {OUTPUT_DIR}")
    print(f"  - 原始转录: *_raw.txt")
    print(f"  - 润色文本: *_polished.txt")
    print(f"  - 完整数据: batch_results.json\n")


# ============= 主函数 =============

def main():
    parser = argparse.ArgumentParser(description="批量视频转录+润色")
    parser.add_argument("videos", nargs="+", help="视频文件路径")
    parser.add_argument("--workers", type=int, default=4, help="转录进程数（默认4）")
    parser.add_argument("--concurrency", type=int, default=20, help="润色并发数（默认20）")
    parser.add_argument("--skip-polish", action="store_true", help="跳过润色步骤")
    parser.add_argument("--api-key", type=str, help="DeepSeek API Key")

    args = parser.parse_args()

    # 更新API Key
    if args.api_key:
        global DEEPSEEK_API_KEY
        DEEPSEEK_API_KEY = args.api_key

    # 验证视频文件
    video_paths = []
    for video in args.videos:
        path = Path(video)
        if path.is_file():
            video_paths.append(str(path))
        else:
            print(f"警告: 文件不存在 - {video}")

    if not video_paths:
        print("错误: 没有有效的视频文件")
        return

    print(f"\n{'='*80}")
    print(f"批量视频转录+润色系统")
    print(f"{'='*80}")
    print(f"总视频数: {len(video_paths)}")
    print(f"转录模型: {WHISPER_MODEL}")
    print(f"润色模型: DeepSeek-Chat")
    print(f"{'='*80}\n")

    # 阶段1: 转录
    transcribe_results = batch_transcribe_multiprocess(
        video_paths,
        num_workers=args.workers
    )

    # 阶段2: 润色
    if not args.skip_polish:
        if DEEPSEEK_API_KEY == "your-api-key-here":
            print("警告: 未设置DeepSeek API Key，跳过润色步骤")
            print("使用 --api-key 参数设置，或在代码中修改 DEEPSEEK_API_KEY\n")
            final_results = transcribe_results
        else:
            final_results = asyncio.run(
                batch_polish_async(
                    transcribe_results,
                    concurrency=args.concurrency
                )
            )
    else:
        final_results = transcribe_results

    # 保存结果
    save_results(final_results)

    print(f"{'='*80}")
    print(f"全部完成！")
    print(f"{'='*80}\n")


if __name__ == "__main__":
    main()
